package com.ruoyi.langchain4j.config;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.elasticsearch.ElasticsearchConfigurationScript;
import dev.langchain4j.store.embedding.elasticsearch.ElasticsearchEmbeddingStore;
import org.elasticsearch.client.RestClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * @description:
 * @author: lizhihua
 * @date: 2025/7/13
 */
@Configuration
public class EmbeddingStoreConfig {

    @Autowired
    private RestClient restClient;

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {

        /**
         * ElasticsearchEmbeddingStore 有 2 种实现：
         *
         * 使用近似 kNN 查询 和 ElasticsearchConfigurationKnn 配置类（默认）。
         * 使用 scriptScore 查询 和 ElasticsearchConfigurationScript 配置类。请注意，此实现使用余弦相似度。
         */
        ElasticsearchEmbeddingStore embeddingStore = ElasticsearchEmbeddingStore
                .builder()
                .indexName("xiaozhi")
//                .configuration(ElasticsearchConfigurationScript.builder().build())
                .restClient(restClient)
                .build();
        return embeddingStore;
    }

}
